This pull request introduces significant updates to the microbenchmarking framework, focusing on new model types, and enhanced shape generation options. The changes aim to expand functionality, and more extensive benchmarking configurations.

Enhancements to Model Types and Shape Generation

• Added support for new model types, including ln_linear_<activation> (e.g., sigmoid, relu, gelu) and transformer_block with self-attention and MLP. These are documented in benchmarks/microbenchmarks/README.md.
• Introduced multiple shape generation options (custom, llama, pow2, pow2_extended, sweep) to support diverse matrix shapes for benchmarking. These options are implemented in benchmark_runner.py and documented in the README.

Refactoring and Code Simplification

• Refactored model creation logic by replacing create_model_and_input with create_model_and_input_data, now imported from torchao.testing.model_architectures. This centralizes model definitions and input data generation.
• Removed redundant model definitions (ToyLinearModel, LNLinearSigmoid) from utils.py, consolidating them into torchao.testing.model_architectures.

Future TODO: Refactor Torchao to use model definitions from torchao.testing.model_architectures, #2078

Updates to Configuration

• Expanded benchmark_config.yml to include configurations for new model types and shape generation options, such as llama and pow2.

Documentation Improvements

• Updated README.md to provide detailed descriptions of new model types and shape generation options, ensuring users can easily understand and utilize the new features.

These changes collectively enhance the flexibility, maintainability, and usability of the benchmarking framework.

Sample configuration.yml for inference benchmarks

benchmark_mode: "inference"
quantization_config_recipe_names:
  - "float8dq"
  - "float8wo"
output_dir: "benchmarks/microbenchmarks/results"
model_params:
  - name: "ln_linear_sigmoid_cuda"
    matrix_shapes:
      - name: "custom"
        shapes: [
          [2048, 4096, 1024],
        ]
    high_precision_dtype: "torch.bfloat16"
    use_torch_compile: true
    torch_compile_mode: "max-autotune"
    device: "cuda"
    model_type: "ln_linear_sigmoid"
    enable_profiler: true

  - name: "bf16_transformer_block"
    matrix_shapes:
      - name: "custom"
        shapes: [
          [2048, 4096, 1024],  # For transformer_block, k is the hidden dimension
        ]
    high_precision_dtype: "torch.bfloat16"
    use_torch_compile: true
    torch_compile_mode: "max-autotune"
    device: "cuda"
    model_type: "transformer_block" # TODO: Add a custom model (Figure out how to do this, maybe pass a .py file with model definition)
    enable_profiler: true

  - name: "large_bf16_ln_linear"
    matrix_shapes:
      - name: "llama"  # Example of using LLaMa shapes
      - name: "pow2"  # Example of using power of 2 shapes
        min_power: 10  # 1024
        max_power: 12  # 4096
      - name: "pow2_extended"  # Example of using extended power of 2 shapes
        min_power: 10  # 1024
        max_power: 11  # 2048
      - name: "sweep"  # Example of using sweep shapes (commented out as it generates many shapes)
        min_power: 8   # 256
        max_power: 9   # 512
    high_precision_dtype: "torch.bfloat16"
    use_torch_compile: true
    torch_compile_mode: "max-autotune"
    device: "cuda"
    model_type: "linear"
    enable_profiler: true  # Enable profiling for this model